package com.gel.aiagent.demo.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.stereotype.Component;

import java.util.Collection;
import java.util.List;
/*多查询扩展*/
@Component
public class MultiQueryExpanderDemo {
    private final ChatClient.Builder chatClientBuilder;

    public MultiQueryExpanderDemo(ChatModel dashscopeChatModel) {
        this.chatClientBuilder = ChatClient. builder(dashscopeChatModel);

    }

    public List<Query> expand(String query){
        MultiQueryExpander multiQueryExpander = MultiQueryExpander.builder()
                .chatClientBuilder(chatClientBuilder)
                .numberOfQueries(5)
                .build();
        return multiQueryExpander.expand(new Query( query));
    }
}
